# Inférence
import tensorflow as tf
import tensorflow.contrib.slim as slim

from config import Input_shape, score_threshold, ignore_thresh, scale_num, model_type, num_class, MAX_BOX_NUM
from my_batch_norm import bn_layer_top


def yolo_head(feature_maps, anchors, num_classes,input_shape, calc_loss=False):
    """
    # pred_xy, pred_wh 是满足 label 定义的预测,yolo_head 是将预测结果由论文定义的格式(相对于当前grid的位置)转换为 label 的
    # 格式(相对于全图的位置)

    Convert final layer features to bounding box parameters.
    (Features learned by the convolutional layers ---> a classifier/regressor which makes the detection prediction)
    :param feature_maps: the feature maps learned by the convolutional layers
                         3 scale [None, 13, 13, 255] from yolov3 structure anchors:[116, 90], [156, 198], [373, 326]
                                 [None, 26, 26, 255]                               [30, 61], [62, 45], [59, 119]
                                 [None, 52, 52, 255]                               [10, 13], [16, 30], [33, 23]
    :param anchors: 3 anchors for each scale shape=(3,2)
    :param input_shape: 416,416
    :return: box_xy  [None, 13, 13, 3, 2], 2: x,y center point of BB
             box_wh  [None, 13, 13, 3, 2], 2: w,h
             box_conf  [None, 13, 13, 3, 1], 1: conf
             box_class_pred  [None, 13, 13, 3, 80], 80: prob of each class
    """
    num_anchors = len(anchors)  # 3
    # Reshape to batch, height, width, num_anchors, box_params
    anchors_tensor = tf.cast(anchors, dtype=feature_maps.dtype)
    anchors_tensor = tf.reshape(anchors_tensor, [1, 1, 1, num_anchors, 2])  # shape=[1,1,1,3,2]

    # CREATE A GRID FOR EACH SCALE
    with tf.name_scope('Create_GRID'):
        grid_shape = tf.shape(feature_maps)[1:3]  # height, width ---> grid 13x13 for scale1
        #         (0,0) (1,0) ...                                      grid 26x26 for scale2
        #         (0,1) (1,1) ...                                      grid 52x52 for scale3
        #          ...
        # In YOLO the height index is the inner most iteration.
        grid_y = tf.range(0, grid_shape[0])  # array([0,1,...,11,12])
        grid_x = tf.range(0, grid_shape[1])
        grid_y = tf.reshape(grid_y, [-1, 1, 1, 1])  # shape=([13,  1,  1,  1])
        grid_x = tf.reshape(grid_x, [1, -1, 1, 1])  # [1, 13, 1, 1]
        grid_y = tf.tile(grid_y, [1, grid_shape[1], 1, 1])  # [13, 1, 1, 1] ---> [13, 13, 1, 1]
        grid_x = tf.tile(grid_x, [grid_shape[0], 1, 1, 1])  # [1, 13, 1, 1] ---> [13, 13, 1, 1]
        grid = tf.concat([grid_x, grid_y], axis=-1)  # shape=[13, 13,  1,  2]
        grid = tf.cast(grid, dtype=feature_maps.dtype)  # change type

    # Reshape [None, 13, 13, 33] =>[None, 13, 13, 3, 11]
    feature_maps_reshape = tf.reshape(feature_maps, [-1, grid_shape[0], grid_shape[1], num_anchors, num_classes+7])
    # prediction 的定义按照论文
    with tf.name_scope('top_feature_maps'):
        # top of feature maps is a activation function
        """softmax is used for the multi-class logistic regression: ouput fall into [-1,1] ---> sum(all classes) = 1
        # sigmoid for the the 2-class logistic regression: output fall into [0,1] ---> sum(all classes) >1
        # for the relative width and weight, use the exponential function"""
        # 中心点在本cell内的相对位置
        box_xy = tf.sigmoid(feature_maps_reshape[..., :2], name='x_y')  # [None, 13, 13, 3, 2]
        box_seed =feature_maps_reshape[..., 4:6]  # [None, 13, 13, 3, 2]
        # wh 相对于每个 anchor 的大小
        box_wh = tf.exp(feature_maps_reshape[..., 2:4], name='w_h')  # [None, 13, 13, 3, 2]
        box_confidence = tf.sigmoid(feature_maps_reshape[..., 6:7], name='confidence')  # [None, 13, 13, 3, 1]
        # box_class_probs = tf.sigmoid(feature_maps_reshape[..., 7:], name='class_probs')  # [None, 13, 13, 3, 80]
        box_class_probs = feature_maps_reshape[..., 7:]  # [None, 13, 13, 3, 4]

        # Adjust predictions to each spatial grid point and anchor size.
        # box_xy 换算成在全图中相对位置， 因为label的定义就是全图中相对位置 例如 (0.5+6)/13 就是全图的 1/2 处
        box_xy = (box_xy + grid) / tf.cast(grid_shape[::-1],  # (x,y + grid)/13. ---> in between (0., 1.)
                                           dtype=feature_maps_reshape.dtype)  # [None, 13, 13, 3, 2]
        box_seed = (box_seed + grid) / tf.cast(grid_shape[::-1],  # (x,y + grid)/13. ---> in between (0., 1.)
                                           dtype=feature_maps_reshape.dtype)  # [None, 13, 13, 3, 2]
        # box_wh 换算成和全图的长度比
        box_wh = box_wh * anchors_tensor / tf.cast(input_shape[::-1],  # following to the scale
                                                   dtype=feature_maps_reshape.dtype)  # [None, 13, 13, 3, 2]

    if calc_loss == True:
        return grid, feature_maps_reshape, box_xy, box_wh
    return box_xy, box_wh, box_seed,box_confidence,box_class_probs

def yolo_correct_boxes(box_xy, box_wh, box_seed,input_shape, image_shape):
    """
    由在图中的相对位置，根据图的长宽，换算成绝对位置
    Convert YOLO box predictions to bounding box corners.(get corrected boxes)
    :param box_xy: (None, 13, 13, 3, 2) is box_x,y output of yolo_head()
    :param box_wh: (None, 13, 13, 3, 2) is box_w,h output of yolo_head()
    :param input_shape: 416,416
    :param image_shape: shape of input image (normal: height, width) . tf.placeholder(shape=(2, ))
    :return: box(2 corners) in original image shape (BB corresponding to h and w of image)
                        : 1 (y_min,x_min) left bottom corner
                          1 (y_max,x_max) right top corner
                    ---> (..., (y_min,x_min,y_max,x_max)) (None, 13, 13, 3, 4)
    """
    # Note: YOLO iterates over height index before width index.
    # batch_size = 3 #tf.shape(image_shape)[0]
    # 相对于全图的位置 ： 竖直，水平
    box_yx = box_xy[..., ::-1]  # (None, 13, 13, 3, 2) => ex: , x,y --> y,x
    box_seed = box_seed[..., ::-1]  # (None, 13, 13, 3, 2) => ex: , x,y --> y,x
    # 相对于全图的长度： 高 宽
    box_hw = box_wh[..., ::-1]  # (None, 13, 13, 3, 2) w,h--->h,w
    input_shape = tf.cast(input_shape, dtype=box_yx.dtype)  # ex: (416,416)
    # input_shape = tf.constant(Input_shape, shape=[batch_size, 2], dtype=box_yx.dtype)
    image_shape = tf.cast(image_shape, dtype=box_yx.dtype)  # ex: (640, 480)

    with tf.name_scope('resize_to_scale_correspond'):
        """un image (640, 480) to scale1 (stride 32)(13x13)
        ---> new shape = (13, 10)"""
        constant = (input_shape / image_shape)
        # constant = tf.reshape((input_shape / image_shape), shape=[batch_size, 2])  # 416/640, 416/480
        # min=[]
        min = tf.minimum(constant[0], constant[1])
        # for i in range(batch_size):
        #     #i+=1
        #     x = tf.minimum(constant[i][0], constant[i][1])
        #     min.append(x)
            #min = tf.concat([min, x], axis=0 )
        # min = tf.stack(min)
        # min = tf.reshape(min, shape=[batch_size, 1])
        # min = (min.append(tf.minimum(constant[i][0], constant[i][1])) for i in range(batch_size))
        # min = tf.cast([min[0], min[1], min[2]], dtype=constant.dtype)
        # min = tf.reshape(min, shape=[batch_size, 2])

        new_shape = image_shape * min  # 640*(416/640), 480*(416/640)
        new_shape = tf.round(new_shape)  # lam tron ---> (416, 312)

    offset = (input_shape - new_shape) / (input_shape*2.)  # 0,  (416-312)/2/416=0.125
    scale = input_shape / new_shape  # (1, 416/312)

    with tf.name_scope('return_corners_box'):
        # box in scale
        # 变成(416,416)正方形中的坐标，用填充的方式，而不是crop
        box_yx = (box_yx - offset) * scale  # (x-0)*1, (y-0.125)*416/312
        box_seed = (box_seed - offset) * scale  # (x-0)*1, (y-0.125)*416/312
        
        box_hw *= scale  # h*1, w*1.333
        box_mins = box_yx - (box_hw / 2.)  # (x-0)*1-h*1/2 = y_min, (y-0.125)*(416/312)-w*(416/312)/2 = x_min
        box_maxes = box_yx + (box_hw / 2.)  # (x-0)*1+h*1/2 = y_max, (y-0.125)*(416/312)+w*(416/312)/2 = x_max
        boxes = tf.concat([box_mins[..., 0:1],  # y_min 上
                           box_mins[..., 1:2],  # x_min 左
                            box_maxes[..., 0:1],  # y_max 下
                           box_maxes[..., 1:2]],  # x_max 右
                          axis=-1, name='box_in_scale')
        # Scale boxes back to original image shape,
        # y_min*height = 720*(x-0)*1-720*h*1/2
        # x_min*width,
        # y_max*height,
        # x_max*width
        # 换算到原图(image_shape)中的位置和长度 (None, 13, 13, 3, 4)
        boxes = tf.multiply(boxes,
                            tf.concat([image_shape, image_shape], axis=-1), name='box_in_original_image_shape')
        box_seed=tf.multiply(box_seed,image_shape, name='seed_in_original_image_shape')
    return boxes,box_seed


def yolo_boxes_and_scores(feats, anchors, input_shape, num_class,image_shape):
    """
    Process Conv layer output
    :param feats: [None, 13, 13, 255] is output of build_networks() (from network_function scripts)
    :param anchors: for yolo_head()
    :param input_shape: see yolo_head(), yolo_correct_boxes()
    :param image_shape: tensor targets for filtered bounding boxes. tf.placeholder(shape=(2, ))
    :return: boxes: [None*13*13*3, 4], predicted BBs with 4: cordoning of one BB (y_min,x_min, y_max,x_max)
             box_scores: [None*13*13*3, 80], 80: score= confidence * class_probability
    """
    # 换算成相对于416的位置和大小
    box_xy, box_wh, box_seed,box_confidence,box_class_probs = yolo_head(feats, anchors, num_class,input_shape)
    # box_yx = box_xy[..., ::-1]  # (None, 13, 13, 3, 2) => ex: , x,y --> y,x
    # # 相对于全图的长度
    # box_hw = box_wh[..., ::-1]  # (None, 13, 13, 3, 2) w,h--->h,w
    rel_boxes=tf.concat([box_xy,box_wh],axis=-1) # 左上宽高
    rel_boxes = tf.reshape(rel_boxes, [-1, 4], name='rel_boxes')  # shape = (None*13*13*3, 4)

    # 根据原图大小，换算成绝对位置和绝对长度 (None, 13, 13, 3, 4)  左上右下===>上左下右 为了做 nms
    boxes,box_seed = yolo_correct_boxes(box_xy, box_wh,box_seed, input_shape,
                               image_shape)  # shape = (None, 13, 13, 3, 4), 4:(y_min,x_min,y_max,x_max)
    boxes = tf.reshape(boxes, [-1, 4], name='boxes')  # shape = (None*13*13*3, 4)
    seeds = tf.reshape(box_seed, [-1, 2], name='boxes')  # shape = (None*13*13*3, 4)
    box_class_probs = tf.reshape(box_class_probs, [-1, num_class], name='boxes')  # shape = (None*13*13*3, 4)
    
    with tf.name_scope('box_scores'):
        box_scores = tf.reshape(box_confidence, [-1])  # (None*13*13*3)
    class_output = tf.argmax(box_class_probs, axis=-1)

    return boxes, rel_boxes,box_scores,seeds,class_output


def order_matrix_net(input, reuse, is_training):
    '''
    regress the order relation of detected box
    :param input: [None,Max_box,4]
    :param reuse:
    :param is_training:
    :return:
    '''
    input.set_shape([None, MAX_BOX_NUM, 4])
    with tf.variable_scope("order_loss_net",reuse=reuse):
        with slim.arg_scope([bn_layer_top],
                            is_training=is_training):

            net=input
            net1=tf.tile(tf.expand_dims(net,axis=2), [1, 1, MAX_BOX_NUM, 1]) #[m,Max_box,Max_box,4]
            # test_output1=net1[0,...,0]

            net2=tf.tile(tf.expand_dims(net,axis=1), [1, MAX_BOX_NUM, 1, 1]) #[m,Max_box,Max_box,4]
            # test_output2=net2[0,...,0]

            net=tf.concat([net1,net2],axis=-1) #[m,Max_box,Max_box,8]
            # test_output=net[0,...,4]
            net=slim.fully_connected(net,64,normalizer_fn=bn_layer_top) #[m,Max_box,Max_box,32]
            net=slim.fully_connected(net,32,normalizer_fn=bn_layer_top) #[m,Max_box,Max_box,32]
            net=slim.fully_connected(net,16,normalizer_fn=bn_layer_top) #[m,Max_box,Max_box,32]
            net=slim.fully_connected(net,1,activation_fn=None) #[m,Max_box,Max_box,1]
            net=tf.squeeze(net,axis=-1) #[m,Max_box,Max_box]
            return net #[m,Max_box,Max_box]
            

def predict(yolo_outputs, anchors,  image_shape, max_boxes=30, score_threshold=0.5, iou_threshold=ignore_thresh):
    """
    Evaluate YOLO model on given input and return filtered boxes
    :param yolo_outputs:[(N,13,13,3,6), (N,26,26,3,6)]
    :param anchors: [6,2]
    :param image_shape: see yolo_boxes_and_scores() 输入图片的真实尺寸
    :param max_boxes: a scalar integer who present the maximum number of boxes to be selected by non max suppression
    :param score_threshold: score_threshold=.6
    :param iou_threshold: iou_threshold=.5
    :return:
    """
    # input_shape = tf.shape(yolo_outputs[0])[1:3] * 32  # scale1 13*32=416 [416,416]
    boxes = []
    seeds = []
    rel_boxes=[]
    box_scores = []
    class_output=[]
    input_shape = (Input_shape, Input_shape)
    if model_type=='N':
        anchor_mask = [[6, 7, 8], [3, 4, 5], [0, 1, 2]]
    else:
        anchor_mask = [[3, 4, 5], [0, 1, 2]]



    # 转换真实大小，求nms
    for scale in range(scale_num):  # 3 scale
        name = 'predict' + str(scale+1)
        with tf.name_scope(name):
            # predict中输入batch_size为1
            # 根据原图大小，换算成绝对位置和绝对长度 (1*13*13*3, 4)
            # _boxes 上左下右，绝对长度，用来nms;   _rel_boxes 左上宽高,相对长度，用来排序
            _boxes, _rel_boxes,_box_scores,_seeds,_class_output = yolo_boxes_and_scores(yolo_outputs[scale],
                                                        anchors[anchor_mask[scale]],
                                                        input_shape,
                                                        num_class,
                                                        image_shape)

            boxes.append(_boxes)  # list(3 array): [3, 1*13*13*3, 4]
            box_scores.append(_box_scores)  # list(3 array): [3, 1*13*13*3]
            rel_boxes.append(_rel_boxes)
            seeds.append(_seeds)
            class_output.append(_class_output)

    boxes = tf.concat(boxes, axis=0)  # [3 *1*13*13*3, 4]
    rel_boxes = tf.concat(rel_boxes, axis=0)  # [3 *1*13*13*3, 4]
    box_scores = tf.concat(box_scores, axis=0)  # [3 *1*13*13*3]
    seeds = tf.concat(seeds, axis=0)  # [3 *1*13*13*3,2]
    class_output = tf.concat(class_output, axis=0)  # [3 *1*13*13*3,1]

    mask = box_scores >= score_threshold  # False & True in [3*1*13*13*3] based on box_scores
    # maximum number of boxes to be selected by non max suppression
    max_boxes_tensor = tf.constant(max_boxes, dtype='int32', name='max_boxes')


    # Class 类分数大于 score_threshold 的boxes
    detected_boxes = tf.boolean_mask(boxes, mask)  # obj:[3 *13*13*3, 4], mask:[3 *13*13*3] ---> [num_box(True), 4], each class: keep boxes who have (box_scores >= score_threshold)
    detected_rel_boxes = tf.boolean_mask(rel_boxes, mask)  # obj:[3 *13*13*3, 4], mask:[3 *13*13*3] ---> [num_box(True), 4], each class: keep boxes who have (box_scores >= score_threshold)
    detected_seeds = tf.boolean_mask(seeds, mask)  # obj:[3 *13*13*3, 4], mask:[3 *13*13*3] ---> [num_box(True), 4], each class: keep boxes who have (box_scores >= score_threshold)
    detected_classes = tf.boolean_mask(class_output, mask)  # obj:[3*13*13*3, 4], mask:[3*13*13*3] ---> [num_box(True), 4], each class: keep boxes who have (box_scores >= score_threshold)
    detected_scores = tf.boolean_mask(box_scores, mask)  # [3*13*13*3] [3*13*13*3] ---> [num_box(True)]
    # [M=num_box(True)]
    nms_index = tf.image.non_max_suppression(detected_boxes,  # [num_box(True), 4]
                                             detected_scores,  # [num_box(True)]
                                             max_boxes_tensor,  # 20
                                             iou_threshold=iou_threshold,
                                             name='non_max_suppression')  # return an integer tensor of indices has the shape [M],
    # M <= 20， num of boxes belong to this class
    # 数据制作的标准： 左上右下
    # [M, 4]   上左下右
    detected_boxes = tf.gather(detected_boxes,
                            nms_index, name='detect_boxes')  # Take the elements of indices (nms_index) in the detected_boxes. [M, 4]
    # [M, 4]   左上宽高
    detected_rel_boxes = tf.gather(detected_rel_boxes,
                            nms_index, name='detect_rel_boxes')  # Take the elements of indices (nms_index) in the detected_boxes. [M, 4]
    # [M, 1]
    detected_scores = tf.gather(detected_scores, nms_index, name='Box_score')  # [M, 1]
    # [M, 1]
    detected_seeds = tf.gather(detected_seeds, nms_index, name='Box_score')  # [M, 2]
    # [M, 1]
    detected_classes = tf.gather(detected_classes, nms_index, name='Box_score')  # [M, 1]
    
    is_training=tf.constant(False,tf.bool)
    with tf.name_scope("Loss_and_Detect1"):
        num_true_sample=tf.shape(detected_rel_boxes)[0]
        mask_paddings=[[0, MAX_BOX_NUM - num_true_sample], [0, 0]] #TODO 18
        # [Max_box, 4]
        detected_rel_boxes = tf.pad(detected_rel_boxes, mask_paddings, 'CONSTANT', constant_values=0)
        # [1.,Max_box, 4] #TODO bug: during inference, the matrix is a constant
        order_matrix=order_matrix_net(tf.expand_dims(detected_rel_boxes,axis=0),reuse=False,is_training=is_training)
        order_matrix=tf.squeeze(order_matrix,axis=0)
        order_matrix=order_matrix[:num_true_sample,:num_true_sample]
        order_matrix = tf.sigmoid(order_matrix)
        order_matrix = tf.cast(order_matrix>0.5,tf.float32)
        
        order_vec=tf.reduce_sum(order_matrix,axis=-1)
        # 按照 order 升序排列
        _, index = tf.nn.top_k(order_vec, k=num_true_sample, sorted=True)
        #  [num_true_sample,4]
        detected_boxes = tf.gather(detected_boxes, index[::-1], axis=0)
        detected_scores= tf.gather(detected_scores, index[::-1], axis=0)
        detected_seeds= tf.gather(detected_seeds, index[::-1], axis=0)
        detected_classes= tf.gather(detected_classes, index[::-1], axis=0)

    # [M, 4] 上左下右
    return detected_boxes,detected_scores,detected_seeds,detected_classes


